from sklearn.model_selection import train_test_split
from sklearn.datasets import load_iris

# 加载示例数据集 (鸢尾花数据集)
iris = load_iris()
X, y = iris.data, iris.target

# 划分数据集，test_size=0.2 表示测试集占总数据的 20%，random_state 设置随机种子保证可重复性
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 再次划分训练集，得到训练集和验证集，validation_size=0.25 表示验证集占训练集的 25%
X_train, X_val, y_train, y_val = train_test_split(X_train, y_train, test_size=0.25, random_state=42) # 0.25 x 0.8 = 0.2 of total data

print("训练集特征 X_train shape:", X_train.shape)
print("训练集标签 y_train shape:", y_train.shape)
print("验证集特征 X_val shape:", X_val.shape)
print("验证集标签 y_val shape:", y_val.shape)
print("测试集特征 X_test shape:", X_test.shape)
print("测试集标签 y_test shape:", y_test.shape)